#include "main.h"

bool FFMPEGCodec::registered = false;

FFMPEGCodec::FFMPEGCodec(CodecType type, bool x264rgb)
    :_x264rgb(x264rgb)
{
    if (!registered) {
        avcodec_register_all();
        registered = true;
    }

    if(type == ENCODER){
        initEncoder();
    }else if(type == DECODER){
        initDecoder();
    }
}

void FFMPEGCodec::initEncoder()
{
    AVCodec *pCodec;
    if(_x264rgb)
        pCodec = avcodec_find_encoder_by_name("libx264rgb");
    else
        pCodec = avcodec_find_encoder_by_name("libx264");
    
    if (!pCodec) {
        LOG("Codec encoder not found\n");
        exit(1);
    }

    pCodecCtx = avcodec_alloc_context3(pCodec);
    if (!pCodecCtx) {
        LOG("Could not allocate video codec context\n");
        exit(1);
    }

    av_opt_set(pCodecCtx->priv_data, "preset", gConfig->ffmpegOption.c_str(), 0);
    av_opt_set(pCodecCtx->priv_data, "tune", "zerolatency", 0); 

    pCodecCtx->width = gConfig->totalWidth;
    pCodecCtx->height = gConfig->totalHeight;
	pCodecCtx->time_base=(AVRational){1, 60};
    pCodecCtx->has_b_frames = 0;
    if(_x264rgb)
        pCodecCtx->pix_fmt = AV_PIX_FMT_RGB24;
    else
        pCodecCtx->pix_fmt = AV_PIX_FMT_YUV420P;
    

    if (avcodec_open2(pCodecCtx, pCodec, NULL) < 0) {
        LOG("Could not open codec\n");
        exit(1);
    }

    rgbFrame = av_frame_alloc();
    if (!rgbFrame) {
        LOG("Failed to allocate rgb frame\n");
        exit(1);
    }
    rgbFrame->format = AV_PIX_FMT_RGB24;
    rgbFrame->width  = pCodecCtx->width;
    rgbFrame->height = pCodecCtx->height;
 
    int ret = av_image_alloc(rgbFrame->data, rgbFrame->linesize, rgbFrame->width, rgbFrame->height,
                         AV_PIX_FMT_RGB24, 32);
    if (ret < 0) {
        LOG("Failed to allocate raw picture buffer\n");
        exit(1);
    }

    if(!_x264rgb){
        //used to convert RGB image to YUV frame.
        convertCtx = sws_getContext(gConfig->totalWidth, gConfig->totalHeight, AV_PIX_FMT_RGB24, 
        gConfig->totalHeight, gConfig->totalHeight, AV_PIX_FMT_YUV420P, SWS_FAST_BILINEAR, NULL, NULL, NULL); 
        if(convertCtx == NULL){
            LOG("Failed to get SwsContext\n");
            exit(1);
        }

        yuvFrame = av_frame_alloc();
        if (!yuvFrame) {
            LOG("Failed to allocate yuv frame\n");
            exit(1);
        }
        yuvFrame->format = AV_PIX_FMT_YUV420P;
        yuvFrame->width  = pCodecCtx->width;
        yuvFrame->height = pCodecCtx->height;
    
        ret = av_image_alloc(yuvFrame->data, yuvFrame->linesize, yuvFrame->width, yuvFrame->height,
                            AV_PIX_FMT_YUV420P, 16);
        if (ret < 0) {
            LOG("Failed to allocate raw picture buffer\n");
            exit(1);
        }
    }

    pts = 0;
}

void FFMPEGCodec::initDecoder()
{
    AVCodec *pCodec = avcodec_find_decoder(AV_CODEC_ID_H264);
    if (!pCodec) {
        LOG("Codec decoder not found\n");
        exit(1);
    }

    pCodecCtx = avcodec_alloc_context3(pCodec);
    if (!pCodecCtx) {
        LOG("Could not allocate video codec context\n");
        exit(1);
    }

    pCodecCtx->width = gConfig->totalWidth;
    pCodecCtx->height = gConfig->totalHeight;
    pCodecCtx->has_b_frames = 0;

    if (avcodec_open2(pCodecCtx, pCodec, NULL) < 0) {
        LOG("Could not open codec\n");
        exit(1);
    }

    if(_x264rgb){
        //used to convert GBRP frame to RGB image.
        convertCtx = sws_getContext(gConfig->totalWidth, gConfig->totalHeight, AV_PIX_FMT_GBRP, 
        gConfig->totalWidth, gConfig->totalHeight, AV_PIX_FMT_RGB24, SWS_FAST_BILINEAR, NULL, NULL, NULL); 
    } else {
        //used to convert YUV frame to RGB image.
        convertCtx = sws_getContext(gConfig->totalWidth, gConfig->totalHeight, AV_PIX_FMT_YUV420P, 
        gConfig->totalWidth, gConfig->totalHeight, AV_PIX_FMT_RGB24, SWS_FAST_BILINEAR, NULL, NULL, NULL); 
    }

    if(convertCtx == NULL){
        LOG("Failed to get SwsContext\n");
        exit(1);
    }

    //when using x264rgb, it's actually GBRP frame, 
    //just don't want to define another variable
    yuvFrame = av_frame_alloc();
    if (!yuvFrame) {
        LOG("Failed to allocate yuv frame\n");
        exit(1);
    }

    rgbFrame = av_frame_alloc();
    if (!rgbFrame) {
        LOG("Failed to allocate rgb frame\n");
        exit(1);
    }

    rgbFrame->format = AV_PIX_FMT_RGB24;
    rgbFrame->width  = pCodecCtx->width;
    rgbFrame->height = pCodecCtx->height;
 
    int ret = av_image_alloc(rgbFrame->data, rgbFrame->linesize, rgbFrame->width, rgbFrame->height,
                         AV_PIX_FMT_RGB24, 32);
    if (ret < 0) {
        LOG("Failed to allocate raw picture buffer\n");
        exit(1);
    }
}   

int FFMPEGCodec::encode(byte* imgData, byte* pktData)
{
    int ret = 0, got_packet = 0;
    AVPacket pkt;
    av_init_packet(&pkt);
    pkt.data = NULL;
    pkt.size = 0;

    memcpy(rgbFrame->data[0], imgData, pCodecCtx->height * pCodecCtx->width * 3);

    // encode into video frame
    if(_x264rgb){
        rgbFrame->pts = pts++;
        ret = avcodec_encode_video2(pCodecCtx, &pkt, rgbFrame, &got_packet);
        if (ret < 0) {
            LOG("Error encoding rgb frame\n");
            return -1;
        }
    } else {
        sws_scale(convertCtx, rgbFrame->data, rgbFrame->linesize, 0, pCodecCtx->height, yuvFrame->data, yuvFrame->linesize);

        yuvFrame->pts = pts++;
        ret = avcodec_encode_video2(pCodecCtx, &pkt, yuvFrame, &got_packet);
        if (ret < 0) {
            LOG("Error encoding yuv frame\n");
            return -1;
        }
    }

    if (got_packet) {
        #ifdef CODEC_LOG
            fprintf(f, "frame: %ld size:%d\n", pts, pkt.size);
        #endif
        memcpy(pktData, pkt.data, pkt.size);
        ret = pkt.size;
    }
    av_free_packet(&pkt);

    return ret;
}


int FFMPEGCodec::decode(byte* pktData, int pktSize, byte* imgData)
{
    int ret = 0, got_packet = 0;
    AVPacket pkt;
    av_init_packet(&pkt);
    pkt.data = pktData;
    pkt.size = pktSize;

    // decode video frame
    ret = avcodec_decode_video2(pCodecCtx, yuvFrame, &got_packet, &pkt);
    if (ret < 0) {
        LOG("Error decoding frame\n");
        return -1;
    }

    sws_scale(convertCtx, yuvFrame->data, yuvFrame->linesize, 0, pCodecCtx->height, rgbFrame->data, rgbFrame->linesize);
    
    if (got_packet) {
        int width = pCodecCtx->width, height = pCodecCtx->height;
        int fsize = rgbFrame->linesize[0] * rgbFrame->height;
        int size = 54 + fsize;

        byte bmp_file_header[14] = { 'B', 'M', 0, 0, 0, 0, 0, 0, 0, 0, 54, 0, 0, 0, };
        byte bmp_info_header[40] = { 40, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 24, 0, };
        byte bmp_pad[3] = { 0, 0, 0 };

        bmp_file_header[2] = (unsigned char)size;
        bmp_file_header[3] = (unsigned char)(size >> 8);
        bmp_file_header[4] = (unsigned char)(size >> 16);
        bmp_file_header[5] = (unsigned char)(size >> 24);

        bmp_info_header[4] = (unsigned char)(width);
        bmp_info_header[5] = (unsigned char)(width >> 8);
        bmp_info_header[6] = (unsigned char)(width >> 16);
        bmp_info_header[7] = (unsigned char)(width >> 24);
        bmp_info_header[8] = (unsigned char)(height);
        bmp_info_header[9] = (unsigned char)(height >> 8);
        bmp_info_header[10] = (unsigned char)(height >> 16);
        bmp_info_header[11] = (unsigned char)(height >> 24);

        memcpy(imgData, bmp_file_header, 14);
        memcpy(imgData + 14, bmp_info_header, 40);
        memcpy(imgData + 54, rgbFrame->data[0], fsize);
        ret = size;
    }
    av_free_packet(&pkt);

    return ret;
}